#include "tftp-server.h"

using namespace std;

int main(int argc, char** argv) {

	setvbuf(stdout, NULL, _IONBF, 0);
	setvbuf(stderr, NULL, _IONBF, 0);

	int status;
	Connection connection;

	status = init(connection);
	if (status == CONNECTION_ERROR) {
		close(connection.socket_descriptor);
		exit(-1);
	}

	status = handler(connection);
	if (status == CONNECTION_ERROR) {
		close(connection.socket_descriptor);
		exit(-1);
	}

	close(connection.socket_descriptor);
	return 0;
}

int init(Connection& connection) {

	int status;

	memset(&connection,0,sizeof(connection));

	int sock = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sock == -1) {
		return CONNECTION_ERROR;
	}

	sockaddr_in address;

	address.sin_family = AF_INET;
	address.sin_port = htons(PORT_NUMBER);

	in_addr inAddr;
	inAddr.s_addr = htonl(INADDR_ANY);
	address.sin_addr = inAddr;

	status = bind(sock, (sockaddr*)&address, sizeof(address));
	if (status == -1) {
		close(sock);
		return CONNECTION_ERROR;
	}

	connection.socket_descriptor = sock;

	return CONNECTION_SUCCESS;
}

int handler(Connection& connection) {

	while(true) {
		socklen_t clientLength = sizeof(connection.client);
		int receivedBytes = recvfrom(
				connection.socket_descriptor,
				connection.buff,BUFFER_SIZE, 0,
				(sockaddr*)&connection.client,
				&clientLength);

		if (receivedBytes == -1) {
			return CONNECTION_ERROR;
		}

		uint16_t opcode = ntohs(*((uint16_t*)connection.buff));

		if (opcode == OPCODE_RRQ || opcode == OPCODE_WRQ) {

			strcpy(connection.filename, connection.buff+2);

			int filenameLength = strlen(connection.filename);

			if (strcmp("octet", connection.buff+2+filenameLength+1)) {
				handleUndefinedError(connection, ERROR_MESSAGE_UNSUPPORTED_MODE);
			} else {
				// we do not bind a port and let the OS handle this for us
				int transferSocket = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
				if (transferSocket == -1) {
					return CONNECTION_ERROR;
				}

				Connection clientConnection = Connection(connection);
				clientConnection.socket_descriptor = transferSocket;

				if (opcode == OPCODE_RRQ) {
					handleRead(clientConnection);
				} else {
					handleWrite(clientConnection);
				}

				// Don't check error on close. we're optimistic
				// and assuming this is a local problem
				close(clientConnection.socket_descriptor);
			}

			memset(&connection.client, 0, sizeof(connection.client));

		} else {
			connection.errorCode = ERROR_CODE_ILLEGAL_TFTP_OPERATION;
			handleError(connection,0);
			continue;
		}
	}
}

int timeoutSelect(Connection& connection) {
	fd_set receiveFds;
	FD_ZERO(&receiveFds);
	timeval timeout;
	timeout.tv_sec = TIMEOUT_SECONDS;
	timeout.tv_usec = 0;
	FD_SET(connection.socket_descriptor, &receiveFds);
	int status = select(connection.socket_descriptor+1, &receiveFds, NULL, NULL, &timeout);
	return status;
}

void handleRead(Connection& connection) {
	connection.blockNumber = 1;
	int blockSize;
	bool retransmit = false;
	int blockRetransmitTries = TIMEOUT_BLOCK_RETRANSMIT_TRIES;
	ifstream readFile;
	readFile.open(connection.filename, ios::binary);

	if (readFile.fail()) {
		connection.errorCode = ERROR_CODE_FILE_NOT_FOUND;
		handleError(connection,0);
		return;
	} else if (!readFile.is_open()) {
		connection.errorCode = ERROR_CODE_ACCESS_VIOLATION;
		handleError(connection,0);
		return;
	}

	// Get length of file:
	readFile.seekg(0, ios::end);
	int length = readFile.tellg();
	readFile.seekg(0, ios::beg);

	do {
		uint16_t opcode = htons(OPCODE_DATA);
		uint16_t blockNumber = htons(connection.blockNumber);

		// Write block number to buffer:
		memcpy(connection.buff, &opcode, sizeof(uint16_t));
		memcpy(connection.buff+2, &blockNumber, sizeof(uint16_t));

		// Read data as a block to buffer:
		blockSize = min(length,512);
		if (retransmit) {
			readFile.seekg(-blockSize, ios::cur);
			retransmit = false;
		}
		readFile.read(connection.buff+4,blockSize);

		socklen_t clientLength = sizeof(connection.client);
		int sentBytes = sendto(
				connection.socket_descriptor,
				connection.buff,blockSize+4, 0,
				(sockaddr*)&connection.client,
				clientLength);

		if (sentBytes == -1) {
			break; // Connection failed, nothing to do
		} else if (sentBytes == 0) {
			break; // Connection disconnected
		} else if (sentBytes < blockSize+4) {
			blockRetransmitTries--;
			if (blockRetransmitTries == 0) {
				break;
			} else {
				continue; // Send again the same block
			}
		}

		clock_t endwait;
		endwait = clock () + TIMEOUT_SECONDS * CLOCKS_PER_SEC ;
		bool timeout = false;

		// Wait for ack:
		do {
	    	int status = timeoutSelect(connection);
	    	if (status == -1) {
	    		goto end_method; // Connection failed, nothing to do
	    	} else if (status == 0) {
	    		timeout = true;
	    		break;
	    	}

			sockaddr_in client;
			memset(&client, 0, sizeof(client));
			clientLength = sizeof(client);
			int receivedBytes = recvfrom(
					connection.socket_descriptor,
					connection.buff,BUFFER_SIZE, 0,
					(sockaddr*)&client,
					&clientLength);

			if (receivedBytes == -1) {
				goto end_method; // Connection failed, nothing to do
			} else if (receivedBytes < OPCODE_ACK_SIZE || client.sin_addr.s_addr != connection.client.sin_addr.s_addr) {
				continue;
			} else if (client.sin_addr.s_addr == connection.client.sin_addr.s_addr && client.sin_port != connection.client.sin_port) {
				connection.errorCode = ERROR_CODE_UNKNOWN_TRANSFER_ID;
				handleError(connection,0);
				continue;
			}

			opcode = ntohs(*((uint16_t*)(connection.buff)));
			blockNumber = ntohs(*((uint16_t*)(connection.buff+2)));
			timeout = clock() > endwait;

			if (opcode == OPCODE_ACK && blockNumber != connection.blockNumber) {
	    		timeout = true;
	    		break;
			}

		} while ((opcode != OPCODE_ACK || blockNumber != connection.blockNumber) && !timeout);

		if (timeout) {
			blockRetransmitTries--;
			if (blockRetransmitTries == 0) {
				handleUndefinedError(connection,ERROR_MESSAGE_BLOCK_RETRANSMISSIONS_EXCEEDED);
				break;
			} else {
				retransmit = true;
				continue; // Send again the same block
			}
		} else {
			// next block:
			connection.blockNumber++;
			length -= blockSize;
			blockRetransmitTries = TIMEOUT_BLOCK_RETRANSMIT_TRIES;
		}

	} while (blockSize == 512);

	end_method: readFile.close();
}

void handleWrite(Connection& connection) {

	int status;
	connection.blockNumber = 0;
	int retransmitRetries = TIMEOUT_BLOCK_RETRANSMIT_TRIES;
	bool transmissionEnded = false;
	bool transmitAck = true;

	// make sure file doesn't already exist
	ifstream checkFile(connection.filename);
	if (checkFile.is_open()) {
		checkFile.close();
		connection.errorCode = ERROR_CODE_FILE_ALREADY_EXISTS;
		handleError(connection,0);
		return;
	}

	ofstream writtenFile;
	writtenFile.open(connection.filename, ios::out | ios::binary);

	// make sure we have write permission
	if (!writtenFile.is_open()) {
		connection.errorCode = ERROR_CODE_ACCESS_VIOLATION;
		handleError(connection,0);
		return;
	}

	while(retransmitRetries > 0) {

		if (transmitAck) {
			status = sendAck(connection);
			if (status == CONNECTION_ERROR) {
				// no reason to continue
				break;
			}
		} else {
			transmitAck = true;
		}

		status = recieveData(connection, transmissionEnded, writtenFile);

    	if (status == CONNECTION_ABORT_CURRENT) {
    		break;
    	} else if (status == CONNECTION_ERROR) {
    		break;
    	} else if (status == CONNECTION_RETRY) {
    		retransmitRetries -= 1;
    	} else if (status == CONNECTION_UNEXPECTED_CLIENT) {
    		transmitAck = false;
    	} else if (status == CONNECTION_TRANSMISSION_ENDED) {
    		// dallying for last ack
    		transmissionEnded = true;
    		retransmitRetries = TIMEOUT_BLOCK_RETRANSMIT_TRIES;
    	} else if (status == CONNECTION_TRANSMISSION_ENDED_ABORT) {
    		break;
    	} else {
    		retransmitRetries = TIMEOUT_BLOCK_RETRANSMIT_TRIES;
    	}

	}

	if (retransmitRetries == 0) {
		handleUndefinedError(connection, ERROR_MESSAGE_BLOCK_RETRANSMISSIONS_EXCEEDED);
	}

	writtenFile.close();

}

int sendAck(Connection& connection) {

	memset(connection.buff, 0, OPCODE_ACK_SIZE);

	uint16_t opcode = htons(OPCODE_ACK);
	uint16_t blockNumber = htons(connection.blockNumber);

	memcpy(connection.buff, &opcode, sizeof(uint16_t));
	memcpy(connection.buff+2, &blockNumber, sizeof(uint16_t));

	int clientLength = sizeof(connection.client);
	int sentBytes = sendto(
			connection.socket_descriptor,
			connection.buff, OPCODE_ACK_SIZE, 0,
			(sockaddr*)&connection.client,
			clientLength);

	if (sentBytes == -1) {
		return CONNECTION_ERROR;
	} else if (sentBytes < OPCODE_ACK_SIZE) {
		return CONNECTION_ERROR;
	}

	return CONNECTION_SUCCESS;
}

int recieveData(Connection& connection, bool transmissionEnded, ofstream& writtenFile) {

	int status = timeoutSelect(connection);

	if (status == -1) {
		return CONNECTION_ERROR;
	} else if (status == 0) {
		if (!transmissionEnded) {
			return CONNECTION_RETRY;
		} else {
			// assuming last ack was recieved by client
			return CONNECTION_TRANSMISSION_ENDED_ABORT;
		}
	}

	sockaddr_in client;
	memset(&client, 0, sizeof(client));

	socklen_t clientLength = sizeof(client);
	int receivedBytes = recvfrom(
			connection.socket_descriptor,
			connection.buff,BUFFER_SIZE, 0,
			(sockaddr*)&client,
			&clientLength);

	if (receivedBytes == -1) {
		return CONNECTION_ERROR;
	} else if (client.sin_addr.s_addr != connection.client.sin_addr.s_addr) {
		return CONNECTION_UNEXPECTED_CLIENT;
	} else if (client.sin_port != connection.client.sin_port) {
		Connection clientConnection = Connection(connection);
		clientConnection.client = client;
		clientConnection.errorCode = ERROR_CODE_UNKNOWN_TRANSFER_ID;
		handleError(clientConnection,0);
		return CONNECTION_RETRY;
	} else if (receivedBytes < 4) {
		handleUndefinedError(connection, ERROR_MESSAGE_MALEFORMED);
		return CONNECTION_ABORT_CURRENT;
	}

	uint16_t opcode = ntohs(*((uint16_t*)connection.buff));

	if (opcode == OPCODE_WRQ && connection.blockNumber == 0) {
		return CONNECTION_RETRY;
	} else if (opcode == OPCODE_ERROR) {
		uint16_t errorCode = ntohs(*((uint16_t*)(connection.buff+2)));
		if (errorCode == ERROR_CODE_UNKNOWN_TRANSFER_ID) {
			return CONNECTION_RETRY;
		} else {
			return CONNECTION_ABORT_CURRENT;
		}
	} else if (opcode == OPCODE_DATA) {
		int blockNumber = ntohs(*((uint16_t*)(connection.buff+2)));
		if (transmissionEnded) {
			return CONNECTION_RETRY;
		} else {
			if (blockNumber != connection.blockNumber + 1) {
				return CONNECTION_RETRY;
			}
			writtenFile.write(connection.buff+4, receivedBytes-4);
			connection.blockNumber += 1;
			if (receivedBytes < 516) {
				return CONNECTION_TRANSMISSION_ENDED;
			} else {
				return receivedBytes;
			}
		}
	} else {
		connection.errorCode = ERROR_CODE_ILLEGAL_TFTP_OPERATION;
		handleError(connection,0);
		return CONNECTION_ABORT_CURRENT;
	}
}

int handleUndefinedError(Connection& connection, string msg) {
    connection.errorCode = ERROR_CODE_NOT_DEFINED;
    memcpy(connection.buff, msg.c_str(), msg.length());
    return handleError(connection, msg.length());
}

int handleError(Connection& connection, int messageLength) {
	uint16_t opcode = htons(OPCODE_ERROR);
	uint16_t errorCode = htons(connection.errorCode);

	int packetLength = OPCODE_REGULAR_ERROR_SIZE;
	if (connection.errorCode == ERROR_CODE_NOT_DEFINED) {
		for (int i = messageLength; i >= 0; connection.buff[i+4] = connection.buff[i], i--); // Shift right 4 bytes
		packetLength += messageLength;
	}

	memcpy(connection.buff, &opcode, sizeof(uint16_t));
	memcpy(connection.buff+2, &errorCode, sizeof(uint16_t));
	connection.buff[packetLength - 1] = 0;

	int clientLength = sizeof(connection.client);
	int sentBytes = sendto(
			connection.socket_descriptor,
			connection.buff, packetLength, 0,
			(sockaddr*)&connection.client,
			clientLength);

	if (sentBytes == -1) {
		return CONNECTION_ERROR;
	} else if (sentBytes < packetLength) {
		return CONNECTION_ERROR;
	}
	return CONNECTION_SUCCESS;
}

